﻿using System.Collections.Generic;
using System.Linq;
using Roslyn.Compilers.CSharp;
using Util;
using SD.Tools.Algorithmia.GeneralDataStructures;

namespace CompleterLib
{
    internal class LazyCompletionsGetter : IEnumerable<KeyValuePair<int, MultiValueDictionary<TypeSymbol, Completion>>>
    {
        private readonly IEnumerator<CompleterReturnValue> xs;
        private readonly SortedDictionary<int, MultiValueDictionary<TypeSymbol, Completion>> data;
        private bool? hasNext;
        private int? lastMin;
        private readonly int? toAdd = null;
        /// <summary>
        /// Assume list is empty if no items have score <100.
        /// </summary>
        private const int MAX_SCORE_GAP = 100;

        public LazyCompletionsGetter(IEnumerable<CompleterReturnValue> xs)
        {
            this.xs = xs.GetEnumerator();
            this.hasNext = null;
            this.data = new SortedDictionary<int, MultiValueDictionary<TypeSymbol, Completion>>();
        }
        private LazyCompletionsGetter(IEnumerable<CompleterReturnValue> xs, int toAdd)
            : this(xs)
        {
            this.toAdd = toAdd;
        }

        private void Add(Completion c)
        {
            int score = c.Score;
            TypeSymbol type = c.Type;
            if (toAdd.HasValue)
            {
                score += toAdd.Value;
                c = c.CloneWithScore(score);
            }
            MultiValueDictionary<TypeSymbol, Completion> d;
            if (!data.TryGetValue(score, out d))
            {
                d = new MultiValueDictionary<TypeSymbol, Completion>();
                data[score] = d;
            }
            d.Add(type, c);
        }
        private int? ComputeNext()
        {
            if (hasNext.HasValue && !hasNext.Value)
            {
                return null;
            }
            else if (lastMin.HasValue
                && lastMin - (data.Count == 0 ? 0 : data.Keys.Max()) > MAX_SCORE_GAP)
            {
                // End early if no results and min score is higher than MIN_FIRST_SCORE.
                hasNext = false;
                return null;
            }
            bool readMinScore = false;
            if (hasNext.HasValue)
            {
                readMinScore = xs.Current is MinFutureScore;
            }
            int minScore = 0;
            if (readMinScore)
            {
                minScore = (xs.Current as MinFutureScore).Score;
                lastMin = minScore;
                hasNext = xs.MoveNext();
            }
            bool skipWhile = false;
            if (readMinScore && hasNext.HasValue && hasNext.Value)
            {
                if (xs.Current is Completion)
                {
                    Add(xs.Current as Completion);
                }
                else
                {
                    skipWhile = true;
                }
            }
            if (!skipWhile)
            {
                while ((hasNext = xs.MoveNext()).Value && !(xs.Current is MinFutureScore))
                {
                    if (xs.Current is Completion)
                    {
                        Add(xs.Current as Completion);
                    }
                }
            }
            return minScore;
        }

        public bool ContainsKey(int key)
        {
            if ((this.FullyComputed || lastMin.HasValue && lastMin > key)
                && data.ContainsKey(key))
            {
                return true;
            }
            else if (data.Keys.Count > 0 && (lastMin.HasValue || this.FullyComputed)
                && lastMin > key)
            {
                return false;
            }
            else
            {
                int? minScore;
                while ((minScore = ComputeNext()).HasValue)
                {
                    if (minScore >= key)
                    {
                        if (!(minScore == 0 && hasNext.Value
                            && (xs.Current as MinFutureScore).Score == 0))
                        {
                            return data.ContainsKey(key);
                        }
                    }
                }
                return false;
            }
        }

        public MultiValueDictionary<TypeSymbol, Completion> this[int key]
        {
            get
            {
                ContainsKey(key); // ensure key has been read if it will be
                return data[key];
            }
        }

        private IEnumerable<KeyValuePair<int, MultiValueDictionary<TypeSymbol, Completion>>> GetAll()
        {
            if (lastMin.HasValue)
            {
                foreach (var kv in data.Where(kv => kv.Key <= lastMin))
                {
                    yield return kv;
                }
            }
            int? lastValue = lastMin.HasValue ? lastMin : null;
            int? minScore;
            while ((minScore = ComputeNext()).HasValue)
            {
                if (!(minScore == 0 && hasNext.Value
                    && (xs.Current as MinFutureScore).Score == 0))
                {
                    foreach (var kv in data.Where(kv => kv.Key <= minScore
                        && (!lastValue.HasValue || kv.Key > lastValue.Value)))
                    {
                        yield return kv;
                    }
                    lastValue = minScore;
                }
            }
            foreach (var kv in data.Where(kv =>
                !lastValue.HasValue || kv.Key > lastValue.Value))
            {
                yield return kv;
            }
        }

        public IEnumerator<KeyValuePair<int, MultiValueDictionary<TypeSymbol, Completion>>> GetEnumerator()
        {
            return GetAll().GetEnumerator();
        }
        public IEnumerable<KeyValuePair<int, SortedSet<Completion>>> EnumerateSets()
        {
            foreach (var entry in GetAll())
            {
                yield return new KeyValuePair<int, SortedSet<Completion>>(
                    entry.Key, entry.Value.Values.Union());
            }
        }
        public IEnumerator<KeyValuePair<int, SortedSet<Completion>>> GetSetEnumerator()
        {
            return EnumerateSets().GetEnumerator();
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return GetAll().GetEnumerator();
        }

        public IEnumerable<Completion> EnumerateCompletions()
        {
            foreach (var kv in this)
            {
                foreach (var comps in kv.Value)
                {
                    foreach (var comp in comps.Value.OrderBy(c => c.CompletionText))
                    {
                        yield return comp;
                    }
                }
            }
        }
        public IEnumerable<CompleterReturnValue> EnumerateCompleterRVs()
        {
            foreach (var kv in this)
            {
                yield return new MinFutureScore(kv.Key);
                foreach (var comps in kv.Value)
                {
                    foreach (var comp in comps.Value.OrderBy(c => c.CompletionText))
                    {
                        yield return comp;
                    }
                }
            }
        }

        public bool FullyComputed { get { return hasNext.HasValue && !(hasNext.Value); } }
        public int? MaxMinScore { get { return !FullyComputed ? (int?)null : data.Keys.Count > 0 ? data.Keys.Max() : 0; } }
        public int MinCount { get { return data.Sum(kv => kv.Value.Sum(cs => cs.Value.Count)); } }

        internal static IEnumerable<CompleterReturnValue> AddToScore(IEnumerable<CompleterReturnValue> xs, int toAdd)
        {
            return new LazyCompletionsGetter(xs, toAdd).EnumerateCompleterRVs();
        }
    }

    internal static class LazyCompletionsGetterUtil
    {
        public static IEnumerable<KeyValuePair<int, List<List<SortedSet<Completion>>>>>
            EnumerateSumSets(this List<LazyCompletionsGetter> xs)
        {
            for (int i = 0; xs.Any(x => !x.FullyComputed) || i <= xs.Sum(x => x.MaxMinScore); i++)
            {
                var starts = xs
                    .Select(x => x.EnumerateSets().TakeWhile(kv => kv.Key <= i)
                        .Select(y => y.Singleton()));
                var iStarts = starts.CartesianProduct().Where(ys => ys.Sum(x => x.Key) == i);
                List<List<SortedSet<Completion>>> res = iStarts.Select(ys => ys.Select(y => y.Value).ToList()).ToList();
                yield return new KeyValuePair<int, List<List<SortedSet<Completion>>>>(i, res);
            }
        }
        public static IEnumerable<KeyValuePair<int, List<List<MultiValueDictionary<TypeSymbol, Completion>>>>>
            EnumerateSums(this List<LazyCompletionsGetter> xs)
        {
            for (int i = 0; xs.Any(x => !x.FullyComputed) || i <= xs.Sum(x => x.MaxMinScore); i++)
            {
                var starts = xs
                    .Select(x => x.TakeWhile(kv => kv.Key <= i)
                        .Select(y => y.Singleton()));
                var iStarts = starts.CartesianProduct().Where(ys => ys.Sum(x => x.Key) == i);
                List<List<MultiValueDictionary<TypeSymbol, Completion>>> res = iStarts.Select(ys => ys.Select(y => y.Value).ToList()).ToList();
                yield return new KeyValuePair<int, List<List<MultiValueDictionary<TypeSymbol, Completion>>>>(i, res);
            }
        }
    }
}
